
import cv2
import numpy as np
import matplotlib.pyplot as plt

def extract_foreground_background(image_path, mask_path):
    # 读取图像和掩码
    image = cv2.imread(image_path)
    mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)

    # 确保掩码是二值的
    _, binary_mask = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)
    print(binary_mask)
    # 提取前景（mask=255）和背景（mask=0）
    foreground = cv2.bitwise_and(image, image, mask=binary_mask)
    background = cv2.bitwise_and(image, image, mask=cv2.bitwise_not(binary_mask))

    return foreground, background

if __name__ == '__main__':
    # 使用示例
    image_path = "/home/wc/disk1/datasets/COLUMBIA/Columbia/4cam_splc/canong3_canonxt_sub_01.tif"
    mask_path = "/home/wc/disk1/datasets/COLUMBIA/Columbia/4cam_splc/edgemask/canong3_canonxt_sub_01_edgemask.jpg"

    foreground, background = extract_foreground_background(image_path, mask_path)

    # 使用 matplotlib 可视化
    plt.figure(figsize=(12, 6))

    plt.subplot(1, 3, 1)
    plt.title('Original Image')
    plt.imshow(cv2.cvtColor(cv2.imread(image_path), cv2.COLOR_BGR2RGB))
    plt.axis('off')

    plt.subplot(1, 3, 2)
    plt.title('Foreground')
    plt.imshow(cv2.cvtColor(foreground, cv2.COLOR_BGR2RGB))
    plt.axis('off')

    plt.subplot(1, 3, 3)
    plt.title('Background')
    plt.imshow(cv2.cvtColor(background, cv2.COLOR_BGR2RGB))
    plt.axis('off')

    plt.tight_layout()
    plt.show()